package com.alibaba.datax.plugin.rdbms.util.extend;

import com.alibaba.datax.common.element.*;
import com.alibaba.datax.common.exception.DataXException;
import com.alibaba.datax.common.util.Configuration;
import com.alibaba.datax.plugin.rdbms.reader.Key;
import com.alibaba.datax.plugin.rdbms.util.DBUtil;
import com.alibaba.datax.plugin.rdbms.util.DBUtilErrorCode;
import com.alibaba.druid.pool.DruidDataSourceFactory;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.sql.DataSource;
import java.io.UnsupportedEncodingException;
import java.sql.*;
import java.util.Map;

/**
 * @author: luozehuan
 * @date: 2022/04/07
 **/
public class AddColumnRecordPreHandler implements RecordPreHandler{

    private final Logger logger = LoggerFactory.getLogger(this.getClass());

    /**
     * 关联补充cust_no的数据源
     */
    private DataSource dataSource;

    /**
     * 是否
     */
    private String isValid;

    /**
     * 处理的sql
     */
    private String handleSql;

    /**
     * 数据库相关配置
     */
    private Map<String,Object> datasourceCfg;

    /**
     * 解码使用
     */
    private String mandatoryEncoding;

    public AddColumnRecordPreHandler(Configuration configuration) {
        isValid = configuration.getString(Key.HANDLE_IS_VALID,"N");
        handleSql = configuration.getString(Key.HANDLE_SQL);
        datasourceCfg = configuration.getMap(Key.HANDLE_DATASOURCE_CFG);
        mandatoryEncoding = configuration.getString(Key.MANDATORY_ENCODING,"");
    }

    @Override
    public boolean isAvailable() {
        return "Y".equals(isValid);
    }

    @Override
    public void beforeHandle(Configuration configuration) {
        try {
            dataSource = DruidDataSourceFactory.createDataSource(datasourceCfg);
        } catch (Exception exception) {
            logger.error("datax extend: beforeHandle error:{}",exception);
        }
    }

    @Override
    public void handle(Record record) {
        try {
            // 获取数据库连接
            Connection connection = dataSource.getConnection();
            // 执行sql，获取需要的列
            // select cust_no from table t where t.loan_no = :loan_no
            ResultSet resultSet = DBUtil.query(connection, handleSql);
            // 将增加的列添加到record的末尾
            handleRecord(resultSet,record,mandatoryEncoding);
        } catch (Exception exception) {
            logger.error("datax extend: handle error: {}",exception);
        }
    }

    private void handleRecord(ResultSet rs,Record record,String mandatoryEncoding){
        try {
            ResultSetMetaData metaData = rs.getMetaData();
            int columnNumber = metaData.getColumnCount();
            for (int i = 1; i <= columnNumber; i++) {
                switch (metaData.getColumnType(i)) {
                    case Types.CHAR:
                    case Types.NCHAR:
                    case Types.VARCHAR:
                    case Types.LONGVARCHAR:
                    case Types.NVARCHAR:
                    case Types.LONGNVARCHAR:
                        String rawData;
                        if(StringUtils.isBlank(mandatoryEncoding)){
                            rawData = rs.getString(i);
                        }else{
                            // charsetName : mandatoryEncoding ; 通过 charsetName 解码 byte[]
                            rawData = new String((rs.getBytes(i) == null ? Consts.EMPTY_CHAR_ARRAY :
                                    rs.getBytes(i)), mandatoryEncoding);
                        }
                        record.addColumn(new StringColumn(rawData));
                        break;

                    case Types.CLOB:
                    case Types.NCLOB:
                        record.addColumn(new StringColumn(rs.getString(i)));
                        break;

                    case Types.SMALLINT:
                    case Types.TINYINT:
                    case Types.INTEGER:
                    case Types.BIGINT:
                        record.addColumn(new LongColumn(rs.getString(i)));
                        break;

                    case Types.NUMERIC:
                    case Types.DECIMAL:
                        record.addColumn(new DoubleColumn(rs.getString(i)));
                        break;

                    case Types.FLOAT:
                    case Types.REAL:
                    case Types.DOUBLE:
                        record.addColumn(new DoubleColumn(rs.getString(i)));
                        break;

                    case Types.TIME:
                        record.addColumn(new DateColumn(rs.getTime(i)));
                        break;

                    // for mysql bug, see http://bugs.mysql.com/bug.php?id=35115
                    case Types.DATE:
                        if (metaData.getColumnTypeName(i).equalsIgnoreCase("year")) {
                            record.addColumn(new LongColumn(rs.getInt(i)));
                        } else {
                            record.addColumn(new DateColumn(rs.getDate(i)));
                        }
                        break;

                    case Types.TIMESTAMP:
                        record.addColumn(new DateColumn(rs.getTimestamp(i)));
                        break;

                    case Types.BINARY:
                    case Types.VARBINARY:
                    case Types.BLOB:
                    case Types.LONGVARBINARY:
                        record.addColumn(new BytesColumn(rs.getBytes(i)));
                        break;

                    // warn: bit(1) -> Types.BIT 可使用BoolColumn
                    // warn: bit(>1) -> Types.VARBINARY 可使用BytesColumn
                    case Types.BOOLEAN:
                    case Types.BIT:
                        record.addColumn(new BoolColumn(rs.getBoolean(i)));
                        break;

                    case Types.NULL:
                        String stringData = null;
                        if(rs.getObject(i) != null) {
                            stringData = rs.getObject(i).toString();
                        }
                        record.addColumn(new StringColumn(stringData));
                        break;

                    default:
                        throw DataXException
                                .asDataXException(
                                        DBUtilErrorCode.UNSUPPORTED_TYPE,
                                        String.format(
                                                "您的配置文件中的列配置信息有误. 因为DataX 不支持数据库读取这种字段类型. 字段名:[%s], 字段名称:[%s], 字段Java类型:[%s]. 请尝试使用数据库函数将其转换datax支持的类型 或者不同步该字段 .",
                                                metaData.getColumnName(i),
                                                metaData.getColumnType(i),
                                                metaData.getColumnClassName(i)));
                }
            }
        } catch (Exception exception) {
            logger.info("datax extend: handleRecord error:{}",exception);
        }
    }
}
